package com.wxit.algorithm.tree;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * 二叉搜索树
 */
public class BSTTree1 {

    BSTNode root;

    static class BSTNode {
        int key;
        Object value;
        BSTNode left;
        BSTNode right;


        public BSTNode(int key) {
            this.key = key;
        }

        public BSTNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public BSTNode(int key, Object value, BSTNode left, BSTNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }

    /**
     * 查找关键字对应的值
     * @param key
     * @return
     */
    public Object get(int key) {
        return doGet(root, key);
    }

    public Object doGet(BSTNode node, int key) {
        if (node == null) {
            return null;
        }
        if (key < node.key) {
            return doGet(node.left, key); // 向左找
        } else if (node.key < key) {
            return doGet(node.right, key); //向又找
        } else {
            return node.value;
        }
    }

    //非递归
    public Object noRecursionGet(int key) {
        BSTNode node = root;
        while (node != null) {
            if (key < node.key) {
                node = node.left;
            } else if (node.key < key) {
                node = node.right;
            } else {
                return node.value;
            }
        }
        return null;
    }

    /**
     * 查找最小关键字对应的值
     * @return
     */
    public Object min() {
        return min(root);
    }

    public Object min(BSTNode node) {
        if (node == null) {
            return null;
        }
        BSTNode p = node;
        while (p.left != null) {
            p = p.left;
        }
        return p.value;
    }

    public Object doMin(BSTNode node) {
        if (node == null) {
            return null;
        }
        if (node.left == null) {
            return node.value;
        }
        return doMin(node.left);
    }

    /**
     * 查找最大关键字对应的值
     * @return
     */
    public Object max() {
        return max(root);
    }

    public Object max(BSTNode node) {
        if (node == null) {
            return null;
        }
        BSTNode p = node;
        while (p.right != null) {
            p = p.right;
        }
        return p.value;
    }

    public Object doMax(BSTNode node) {
        if (node == null) {
            return null;
        }
        if (node.right == null) {
            return node.value;
        }
        return doMax(node.right);
    }

    public void put(int key, Object value) {
        BSTNode node = root;
        BSTNode parent = null;
        while (node != null) {
            parent = node;
            if (key < node.key) {
                node = node.left;
            } else if (node.key < key) {
                node = node.right;
            } else {
               // 更新
                node.value = value;
                return;
            }
        }
        if (parent == null) {
            root = new BSTNode(key, value);
            return;
        }
        // 新增
        if (key < parent.key) {
            parent.left = new BSTNode(key, value);
        } else {
            parent.right = new BSTNode(key, value);
        }
    }

    /**
     * 查找关键字的前驱值
     * @param key
     * @return
     */
    public Object successor(int key) {
        BSTNode p = root;
        BSTNode ancestorFromLeft = null;
        while (p != null) {
            if (key < p.key) {
                p = p.left;
            } else if (p.key < key) {
                ancestorFromLeft = p;
                p = p.right;
            } else {
                break;
            }
        }
        //没找到节点
        if (p == null) {
            return null;
        }
        //两种情况
        //如果有左子树，此时前任就是左子树的最大值
        if (p.left != null) {
            return max(p.left);
        }
        //找到节点，节点没有左子树，离他最近，从左而来最近的祖先
        return ancestorFromLeft != null ? ancestorFromLeft.value : null;
    }

    /**
     * 查找关键字的后继值
     * @param key
     * @return
     */
    public Object predecessor(int key) {
        BSTNode p = root;
        //自右而来
        BSTNode ancestorFromRight = null;
        while (p != null) {
            if (key < p.key) {
                ancestorFromRight = p;
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else {
                break;
            }
        }
        //没找到节点
        if (p == null) {
            return null;
        }
        //两种情况
        //如果有右子树，此时前任就是右子树的最小值
        if (p.right != null) {
            return min(p.left);
        }
        //找到节点，节点没有右子树，离他最近，从右而来最近的祖先
        return ancestorFromRight != null ? ancestorFromRight.value : null;
    }

    /**
     * 删除节点
     * @param key
     * @return
     */
    public Object delete(int key) {
        BSTNode p = root;
        BSTNode parent = null;
        while (p != null) {
            if (key < p.key) {
                parent = p;
                p = p.left;
            } else if (p.key < key) {
                parent = p;
                p = p.right;
            } else {
                break;
            }
        }
        if (p == null) {
            return null;
        }
        //删除操作
        if (p.left == null) {
            //情况1
            shift(parent, p, p.right);
        } else if (p.right == null) {
            //情况2
            shift(parent, p, p.left);
        } else {
            //情况4
            //4.1被删除节点的后继
            BSTNode s = p.right;
            BSTNode sParent = p; //后继的父节点
            while (s.left != null) {
                sParent = s;
                s = s.left;
            }
            //后继节点即为s
            if (sParent != p) { //不相邻
                //4.2删除和后继不相邻，处理后继的后事
                shift(sParent, s, s.right); //不可能有左孩子
                s.right = p.right;
            }
            //4.3后继取代被删除的节点
            shift(parent, p, s);
            s.left = p.left;
        }
        return p.value;
    }
    /**
     * 托孤方法
     */
    private void shift(BSTNode parent, BSTNode deleted, BSTNode child) {
        if (parent == null) {
            root = child;
        } else if (deleted == parent.left) {
            parent.left = child;
        } else {
            parent.right = child;
        }
    }

    public Object deleteRecursively(int key) {
        //保存被删除节点的值
        ArrayList<Object> result = new ArrayList<>();
        root = doDelete(root, key, result);
        return result.isEmpty() ? null : result.get(0);
    }

    /**
     * 4
     * /  \
     * 2    6
     * / \    \
     * 1   3    7
     *
     * @param node   删除的起点
     * @param key
     * @param result
     * @return 删除剩下的孩子
     */
    private BSTNode doDelete(BSTNode node, int key, ArrayList<Object> result) {
        if (node == null) {
            return null;
        }
        if (key < node.key) {
            //4->1
            node.left = doDelete(node.left, key, result);
            return node;
        }
        if (node.key < key) {
            node.right = doDelete(node.right, key, result);
            return node;
        }
        result.add(node.value);
        //3种情况
        //只有右孩子
        if (node.left == null) {
            return node.right;
        }
        //只有左孩子
        if (node.right == null) {
            return node.left;
        }
        //有两个孩子
        BSTNode s = node.right;
        while (s.left != null) {
            s = s.left;
        }
        s.right = doDelete(node.right, s.key, new ArrayList<>());
        s.left = node.left;
        return s;
    }


    public List<Object> less(int key) {
        ArrayList<Object> result = new ArrayList<>();
        BSTNode p = root;
        LinkedList<BSTNode> stack = new LinkedList<>();
        while (p != null || !stack.isEmpty()) {
            if (p != null) {
                stack.push(p);
                p = p.left;
            } else {
                BSTNode pop = stack.pop();
                //处理值
                if (pop.key < key) {
                    result.add(pop.value);
                } else {
                    break;
                }
                p = pop.right;
            }
        }
        return result;
    }

    /**
     * 降序遍历
     * @param key
     * @return
     */
    public List<Object> greater(int key) {
        ArrayList<Object> result = new ArrayList<>();
        BSTNode p = root;
        LinkedList<BSTNode> stack = new LinkedList<>();
        while (p != null || !stack.isEmpty()) {
            if (p != null) {
                stack.push(p);
                p = p.right;
            } else {
                BSTNode pop = stack.pop();
                if (pop.key > key) {
                    result.add(pop.value);
                } else {
                    break;
                }
                p = pop.left;
            }
        }
        return result;
    }

    public List<Object> between(int key1, int key2) {
        ArrayList<Object> result = new ArrayList<>();
        BSTNode p = root;
        LinkedList<BSTNode> stack = new LinkedList<>();
        while (p != null || !stack.isEmpty()) {
            if (p != null) {
                stack.push(p);
                p = p.left;
            } else {
                BSTNode pop = stack.pop();
                if (key1 <= pop.key && pop.key <= key2) {
                    result.add(pop.value);
                }
                p = pop.right;
            }
        }
        return result;
    }
}
